"""BLE protocol handler with fragmentation support."""

import json
import logging
import struct
import threading
from typing import Dict, Any, Optional, List, Tuple, Callable
from datetime import datetime, timedelta
from dataclasses import dataclass, field

from src.ble.constants import (
    MTU_DEFAULT,
    MTU_MIN,
    MTU_MAX,
    FRAGMENT_HEADER_SIZE,
    REASSEMBLY_TIMEOUT,
    STATUS_SUCCESS,
    STATUS_ERROR,
    STATUS_INVALID_COMMAND,
    STATUS_FRAGMENTATION_ERROR,
    STATUS_TIMEOUT,
    RESPONSE_PONG,
    RESPONSE_ERROR
)

logger = logging.getLogger(__name__)


@dataclass
class Fragment:
    """Represents a data fragment."""
    sequence: int
    total: int
    data: bytes
    timestamp: datetime = field(default_factory=datetime.now)


class FragmentReassembler:
    """Handles reassembly of fragmented messages."""
    
    def __init__(self, timeout: int = REASSEMBLY_TIMEOUT):
        """
        Initialize fragment reassembler.
        
        Args:
            timeout: Timeout for reassembly in seconds
        """
        self.timeout = timeout
        self.sessions: Dict[str, Dict[int, Fragment]] = {}
        self.lock = threading.RLock()
        self._cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self._cleanup_thread.start()
        self._running = True
        
    def stop(self):
        """Stop the reassembler."""
        self._running = False
        
    def _cleanup_loop(self):
        """Clean up expired reassembly sessions."""
        import time
        while self._running:
            try:
                self._cleanup_expired()
            except Exception as e:
                logger.error(f"Cleanup error: {e}")
            time.sleep(5)
            
    def _cleanup_expired(self):
        """Remove expired fragment sessions."""
        with self.lock:
            now = datetime.now()
            expired = []
            
            for session_id, fragments in self.sessions.items():
                if fragments:
                    first_fragment = min(fragments.values(), key=lambda f: f.timestamp)
                    if (now - first_fragment.timestamp).total_seconds() > self.timeout:
                        expired.append(session_id)
                        
            for session_id in expired:
                logger.warning(f"Fragment session {session_id} expired")
                del self.sessions[session_id]
                
    def add_fragment(self, session_id: str, sequence: int, total: int, data: bytes) -> Optional[bytes]:
        """
        Add a fragment and return complete data if all fragments received.
        
        Args:
            session_id: Session identifier
            sequence: Fragment sequence number (1-based)
            total: Total number of fragments
            data: Fragment data
            
        Returns:
            Complete reassembled data if all fragments received, None otherwise
        """
        with self.lock:
            if session_id not in self.sessions:
                self.sessions[session_id] = {}
                
            fragments = self.sessions[session_id]
            
            # Validate fragment
            if sequence < 1 or sequence > total:
                logger.error(f"Invalid fragment sequence {sequence}/{total}")
                return None
                
            # Add fragment
            fragments[sequence] = Fragment(sequence, total, data)
            
            # Check if all fragments received
            if len(fragments) == total:
                # Reassemble data
                try:
                    complete_data = b""
                    for i in range(1, total + 1):
                        if i not in fragments:
                            logger.error(f"Missing fragment {i}/{total}")
                            return None
                        complete_data += fragments[i].data
                        
                    # Clean up session
                    del self.sessions[session_id]
                    return complete_data
                    
                except Exception as e:
                    logger.error(f"Reassembly error: {e}")
                    del self.sessions[session_id]
                    return None
                    
            return None
            
    def get_session_status(self, session_id: str) -> Dict[str, Any]:
        """Get status of a reassembly session."""
        with self.lock:
            if session_id not in self.sessions:
                return {"exists": False}
                
            fragments = self.sessions[session_id]
            if not fragments:
                return {"exists": True, "fragments": 0}
                
            total = max(f.total for f in fragments.values())
            return {
                "exists": True,
                "fragments": len(fragments),
                "total": total,
                "missing": [i for i in range(1, total + 1) if i not in fragments]
            }


class ProtocolHandler:
    """Handles BLE protocol communication."""
    
    def __init__(self, device_info_callback: Optional[Callable] = None):
        """
        Initialize protocol handler.
        
        Args:
            device_info_callback: Callback to get device information
        """
        self.mtu = MTU_DEFAULT
        self.reassembler = FragmentReassembler()
        self.device_info_callback = device_info_callback
        self.command_handlers: Dict[str, Callable] = {
            "ping": self._handle_ping,
            "get_info": self._handle_get_info,
            "set_mtu": self._handle_set_mtu
        }
        
    def stop(self):
        """Stop the protocol handler."""
        self.reassembler.stop()
        
    def set_mtu(self, mtu: int):
        """
        Set MTU value.
        
        Args:
            mtu: MTU value (will be clamped to valid range)
        """
        self.mtu = max(MTU_MIN, min(mtu, MTU_MAX))
        logger.info(f"MTU set to {self.mtu}")
        
    def fragment_data(self, data: bytes) -> List[bytes]:
        """
        Fragment data into MTU-sized chunks.
        
        Args:
            data: Data to fragment
            
        Returns:
            List of fragments with headers
        """
        max_data_size = self.mtu - FRAGMENT_HEADER_SIZE
        
        # Check if fragmentation is needed
        if len(data) <= max_data_size:
            # No fragmentation needed - send as single fragment
            header = struct.pack(">HH", 1, 1)  # seq=1, total=1
            return [header + data]
            
        # Calculate number of fragments
        fragments = []
        total = (len(data) + max_data_size - 1) // max_data_size
        
        if total > 65535:  # Max fragments limited by 16-bit counter
            raise ValueError(f"Data too large to fragment: {len(data)} bytes")
            
        # Create fragments
        for i in range(total):
            start = i * max_data_size
            end = min(start + max_data_size, len(data))
            fragment_data = data[start:end]
            
            # Add header (sequence, total)
            header = struct.pack(">HH", i + 1, total)
            fragments.append(header + fragment_data)
            
        logger.debug(f"Fragmented {len(data)} bytes into {len(fragments)} fragments")
        return fragments
        
    def defragment_data(self, session_id: str, fragment: bytes) -> Optional[bytes]:
        """
        Process a fragment and return complete data if all fragments received.
        
        Args:
            session_id: Session identifier
            fragment: Fragment with header
            
        Returns:
            Complete data if all fragments received, None otherwise
        """
        if len(fragment) < FRAGMENT_HEADER_SIZE:
            logger.error(f"Fragment too small: {len(fragment)} bytes")
            return None
            
        # Parse header
        sequence, total = struct.unpack(">HH", fragment[:FRAGMENT_HEADER_SIZE])
        data = fragment[FRAGMENT_HEADER_SIZE:]
        
        # Single fragment case
        if total == 1:
            return data
            
        # Multi-fragment case
        return self.reassembler.add_fragment(session_id, sequence, total, data)
        
    def handle_command(self, data: bytes, session_id: str = "default") -> bytes:
        """
        Handle incoming command data.
        
        Args:
            data: Raw command data (may be fragmented)
            session_id: Session identifier for fragmentation
            
        Returns:
            Response data
        """
        try:
            # Handle fragmentation
            complete_data = self.defragment_data(session_id, data)
            if complete_data is None:
                # Waiting for more fragments
                return self._create_response({
                    "status": "waiting_fragments",
                    "session": self.reassembler.get_session_status(session_id)
                })
                
            # Parse JSON command
            try:
                command = json.loads(complete_data.decode('utf-8'))
            except (json.JSONDecodeError, UnicodeDecodeError) as e:
                logger.error(f"Invalid JSON command: {e}")
                return self._create_error_response(STATUS_INVALID_COMMAND, "Invalid JSON format")
                
            # Validate command structure
            if not isinstance(command, dict) or "cmd" not in command:
                return self._create_error_response(STATUS_INVALID_COMMAND, "Missing 'cmd' field")
                
            # Handle command
            cmd_name = command["cmd"]
            if cmd_name in self.command_handlers:
                response = self.command_handlers[cmd_name](command)
            else:
                response = self._create_error_response(
                    STATUS_INVALID_COMMAND, 
                    f"Unknown command: {cmd_name}"
                )
                
            return response
            
        except Exception as e:
            logger.error(f"Command handling error: {e}")
            return self._create_error_response(STATUS_ERROR, str(e))
            
    def _handle_ping(self, command: Dict[str, Any]) -> bytes:
        """Handle ping command."""
        device_id = self._get_device_id()
        return self._create_response({
            "result": RESPONSE_PONG,
            "device_id": device_id,
            "timestamp": datetime.now().isoformat()
        })
        
    def _handle_get_info(self, command: Dict[str, Any]) -> bytes:
        """Handle get_info command."""
        info = {
            "device_id": self._get_device_id(),
            "mtu": self.mtu,
            "protocol_version": "1.0",
            "capabilities": ["wifi_config", "fragmentation"],
            "timestamp": datetime.now().isoformat()
        }
        
        # Add custom device info if callback provided
        if self.device_info_callback:
            try:
                custom_info = self.device_info_callback()
                if isinstance(custom_info, dict):
                    info.update(custom_info)
            except Exception as e:
                logger.error(f"Device info callback error: {e}")
                
        return self._create_response({"result": "success", "info": info})
        
    def _handle_set_mtu(self, command: Dict[str, Any]) -> bytes:
        """Handle set_mtu command."""
        if "mtu" not in command:
            return self._create_error_response(STATUS_INVALID_COMMAND, "Missing 'mtu' field")
            
        try:
            new_mtu = int(command["mtu"])
            if new_mtu < MTU_MIN or new_mtu > MTU_MAX:
                return self._create_error_response(
                    STATUS_INVALID_COMMAND,
                    f"MTU must be between {MTU_MIN} and {MTU_MAX}"
                )
                
            self.set_mtu(new_mtu)
            return self._create_response({
                "result": "success",
                "mtu": self.mtu
            })
            
        except (ValueError, TypeError) as e:
            return self._create_error_response(STATUS_INVALID_COMMAND, "Invalid MTU value")
            
    def _get_device_id(self) -> str:
        """Get device identifier."""
        try:
            from src.utils.dbus_helper import get_device_mac
            mac = get_device_mac()
            if mac:
                return mac.replace(":", "")[-4:]
            return "0000"
        except Exception:
            return "0000"
            
    def _create_response(self, data: Dict[str, Any]) -> bytes:
        """Create JSON response."""
        try:
            json_str = json.dumps(data, separators=(',', ':'))
            return json_str.encode('utf-8')
        except Exception as e:
            logger.error(f"Response creation error: {e}")
            return b'{"status":"error","message":"Response creation failed"}'
            
    def _create_error_response(self, code: int, message: str) -> bytes:
        """Create error response."""
        return self._create_response({
            "status": RESPONSE_ERROR,
            "code": code,
            "message": message
        })